#include "pch.h"
#include "UIThreadTask.h"
#include "..\Common\SRWLock.h"

#define WM_RUN_TASK		WM_USER + 1000
#define TID_RUN_TASK	1

extern HMODULE g_hModule;

CUIThreadTask *CUIThreadTask::m_pThis;

CUIThreadTask::CUIThreadTask()
	: m_hWnd(NULL)
	, m_pSuperProc(NULL)
	, m_bInRunTask(false)
	, m_nTaskDueTime(0)
{
	AssertInUIThread();

	m_pThis = this;
}

CUIThreadTask::~CUIThreadTask()
{
	AssertInUIThread();
	assert(m_hWnd == NULL);
	assert(m_deqTasks.empty());

	m_pThis = nullptr;
}

bool CUIThreadTask::Init()
{
	AssertInUIThread();

	m_hWnd = CreateWindowExW(0, L"Message", NULL, WS_POPUP, 0, 0, 0, 0, HWND_MESSAGE, NULL, g_hModule, NULL);

	WNDPROC pProc = [](HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		return m_pThis->WndProc(hWnd, uMsg, wParam, lParam);
	};

	return m_hWnd && (m_pSuperProc = (WNDPROC)SetWindowLongPtrW(m_hWnd, GWLP_WNDPROC, (LONG_PTR)pProc));
}

void CUIThreadTask::Exit()
{
	AssertInUIThread();

	if (m_hWnd)
	{
		DestroyWindow(m_hWnd);
		m_hWnd = NULL;
	}

	CExclusiveLock lock(m_SRWLock);
	auto deqTasks = std::move(m_deqTasks);
	auto queTimedTasks = std::move(m_queTimedTasks);
	lock.Unlock();

	// Sequentially clean up pending tasks
	while (deqTasks.size())
	{
		deqTasks.pop_front();
	}

	while (queTimedTasks.size())
	{
		queTimedTasks.pop();
	}
}

bool CUIThreadTask::PostTask(ICallbackPtr&& pTask)
{
	CExclusiveLock lock(m_SRWLock);
	m_deqTasks.emplace_back(std::move(pTask));
	lock.Unlock();

	// Go to the UI thread to run the task
	PostMessage(m_hWnd, WM_RUN_TASK, 0, 0);
	return true;
}

bool CUIThreadTask::PostTask(UINT nDelayMs, ICallbackPtr&& pTask)
{
	TimedTask task;
	task.nDueTime = GetTickCount() + nDelayMs;

	UINT nTaskId = pTask->GetId();
	task.pTask = std::move(pTask);

	CExclusiveLock lock(m_SRWLock);
	m_queTimedTasks.emplace(std::move(task));

	// If the new task expires first, go to the UI thread to reset the timer
	if (m_queTimedTasks.top().pTask->GetId() == nTaskId)
	{
		lock.Unlock();
		PostMessage(m_hWnd, WM_RUN_TASK, 0, 0);
	}

	return true;
}

LRESULT CUIThreadTask::WndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	AssertInUIThread();

	switch (uMsg)
	{
	case WM_RUN_TASK:
		RunTask();
		break;

	case WM_TIMER:
		if (wParam == TID_RUN_TASK)
		{
			// Stop the timer before running tasks
			m_nTaskDueTime = 0;
			KillTimer(m_hWnd, TID_RUN_TASK);
			RunTask();
		}
		break;

	default:
		return m_pSuperProc(hWnd, uMsg, wParam, lParam);
	}

	return 0;
}

void CUIThreadTask::RunTask()
{
	AssertInUIThread();

	if (m_bInRunTask)
		return;

	m_bInRunTask = true;

	while (m_hWnd)
	{
		CExclusiveLock lock(m_SRWLock);

		if (m_deqTasks.size())
		{
			auto pTask = std::move(m_deqTasks.front());
			m_deqTasks.pop_front();
			lock.Unlock();

			pTask->Invoke();
			pTask = nullptr;
		}
		else if (m_queTimedTasks.size())
		{
			UINT nDueTime = m_queTimedTasks.top().nDueTime;
			int nElapse = nDueTime - GetTickCount();

			// The first task has not expired
			if (nElapse > 0)
			{
				lock.Unlock();

				if (m_nTaskDueTime != nDueTime)
				{
					m_nTaskDueTime = nDueTime;
					SetTimer(m_hWnd, TID_RUN_TASK, nElapse, NULL);
				}
				break;
			}

			auto pTask = std::move(((TimedTask&)m_queTimedTasks.top()).pTask);
			m_queTimedTasks.pop();
			lock.Unlock();

			// Stop the timer before running the task
			if (m_nTaskDueTime)
			{
				m_nTaskDueTime = 0;
				KillTimer(m_hWnd, TID_RUN_TASK);
			}

			pTask->Invoke();
			pTask = nullptr;
		}
		else
			break;

		// Remove excess messages and handle the quit message
		MSG msg;
		if (PeekMessage(&msg, m_hWnd, WM_RUN_TASK, WM_RUN_TASK, PM_REMOVE) && msg.message == WM_QUIT)
		{
			PostQuitMessage((int)msg.wParam);
			break;
		}
	}

	m_bInRunTask = false;
}
